package com.stone.inventory.sharding.helper;

import com.stone.inventory.sharding.config.properties.TableRuleProperty;
import com.stone.inventory.sharding.id.ShardingIdInjector;
import com.stone.inventory.sharding.id.TableIdGenerator;
import com.stone.inventory.sharding.infra.algorithm.ShardingAlgorithm;
import com.stone.starter.redis.client.RedisShardingClient;
import lombok.Data;
import org.springframework.data.redis.core.StringRedisTemplate;

import javax.persistence.Table;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 分片id助手
 *
 * @author Mr_wenpan@163.com 2021-08-08 21:54:57
 */
@Data
public class ShardingHelper {

    /**
     * 缓存需要分片的实体和该实体对应的逻辑表名，key为实体Class，value为逻辑表名
     */
    private static final Map<Class<?>, String> LOGIC_TABLES = new ConcurrentHashMap<>();

    /**
     * 是否开启分片
     */
    private boolean enableSharding;
    /**
     * 数据库schema
     */
    private String schema;
    /**
     * 缓存redis实例名称
     */
    private String cacheRedisInstance;
    /**
     * 默认redis实例名称
     */
    private String defaultRedisInstance;
    /**
     * redis实际分片实例，key为分片表名，value为该表对应的redis实例名称
     */
    private Map<String, String> redisShardingMapping;
    /**
     * 分片ID注入器集合
     */
    private Map<String, ShardingIdInjector> idInjectorMap;
    /**
     * 表Id生成器集合
     */
    private Map<String, TableIdGenerator> tableIdGeneratorMap;
    /**
     * 分片算法集合，key对应着逻辑表名，value对应着对该表的分片算法
     */
    private Map<String, ShardingAlgorithm<?>> shardingAlgorithmMap;
    /**
     * Redis分片客户端
     */
    private RedisShardingClient redisShardingClient;
    /**
     * 分表规则属性集合
     */
    private Map<String, TableRuleProperty> tableRules;

    public ShardingHelper(boolean enableSharding) {
        this.enableSharding = enableSharding;
        idInjectorMap = new LinkedHashMap<>();
        tableIdGeneratorMap = new LinkedHashMap<>();
        shardingAlgorithmMap = new LinkedHashMap<>();
    }

    /**
     * 注册id注入器
     *
     * @param shardingIdInjector 分片id注入器
     */
    public void register(ShardingIdInjector<?, ?> shardingIdInjector) {
        idInjectorMap.put(shardingIdInjector.logicTableName(), shardingIdInjector);
    }

    /**
     * 注册表id生成器
     *
     * @param tableIdGenerator 表id生成器
     */
    public void register(TableIdGenerator tableIdGenerator) {
        tableIdGeneratorMap.put(tableIdGenerator.getActualTableName(), tableIdGenerator);
    }

    /**
     * 注册分片算法
     *
     * @param shardingAlgorithm 分片算法
     */
    public void register(ShardingAlgorithm<?> shardingAlgorithm) {

        shardingAlgorithmMap.put(shardingAlgorithm.logicTableName(), shardingAlgorithm);
    }

    /**
     * 自动生成主键
     *
     * @param injectionObject 待生成id的实体对象
     */
    public void autoGeneratorKey(Object injectionObject) {

        // 获取逻辑表名
        String logicTable = getEntityLogicTable(injectionObject.getClass());

        // 获取分片字段的值
        ShardingIdInjector shardingIdInjector = idInjectorMap.get(logicTable);
        // 分片字段value
        Object shardingColumnValue = shardingIdInjector.getShardingData(injectionObject);

        // 根据分片字段 得到数据实际要存入的表名称
        ShardingAlgorithm shardingAlgorithm = shardingAlgorithmMap.get(logicTable);
        String actualTable = shardingAlgorithm.getTargetTableByShardingValue(shardingColumnValue);

        // 根据表名称获取该表的id生成器
        Long id = tableIdGeneratorMap.get(actualTable).nextId();
        shardingIdInjector.injection(injectionObject, id);
    }

    /**
     * 获取实际的表名称
     *
     * @param entity 实体对象
     * @return 实际表名称
     */
    public String getActualTableByEntity(Object entity) {

        String logicTable = getEntityLogicTable(entity.getClass());

        // 获取分片字段的值
        ShardingIdInjector shardingIdInjector = idInjectorMap.get(logicTable);
        Object shardingColumnValue = shardingIdInjector.getShardingData(entity);

        // 根据分片字段 得到数据实际要存入的表名称
        ShardingAlgorithm shardingAlgorithm = shardingAlgorithmMap.get(logicTable);
        return shardingAlgorithm.getTargetTableByShardingValue(shardingColumnValue);
    }

    /**
     * 获取redis 分片连接
     *
     * @param logicTable    逻辑表名称
     * @param shardingValue 分片字段值
     * @return redis连接实例
     */
    public StringRedisTemplate getRedisByShardingColumn(Object shardingValue, String logicTable) {

        // 未开启分片使用默认redis实例
        if (!enableSharding || !isSharding(logicTable)) {
            return getRedisByInstanceName(defaultRedisInstance);
        }

        ShardingAlgorithm shardingAlgorithm = shardingAlgorithmMap.get(logicTable);

        // 如果没有分片算法则返回默认的redis实例
        if (null == shardingAlgorithm) {
            return getRedisByInstanceName(defaultRedisInstance);
        }

        // 使用分片算法获取实际的redis实例
        String actualTable = shardingAlgorithm.getTargetTableByShardingValue(shardingValue);
        String actualInstance = redisShardingMapping.get(actualTable);
        return redisShardingClient.getRedisTemplate(actualInstance);
    }

    /**
     * 获取redis分片连接
     *
     * @param logicTable 逻辑表名称
     * @param shardingId 分片Id
     * @return redis 连接实例
     */
    public StringRedisTemplate getRedisByShardingId(String logicTable, long shardingId) {

        // 未开启分片使用默认redis实例
        if (!enableSharding || !isSharding(logicTable)) {
            return getRedisByInstanceName(defaultRedisInstance);
        }
        ShardingAlgorithm shardingAlgorithm = shardingAlgorithmMap.get(logicTable);

        // 如果没有分片算法则返回默认的redis实例
        if (null == shardingAlgorithm) {
            return getRedisByInstanceName(defaultRedisInstance);
        }

        // 使用分片算法获取实际的redis实例
        String actualTable = shardingAlgorithm.getTargetTableByShardingId(shardingId);
        String actualInstance = redisShardingMapping.get(actualTable);
        return redisShardingClient.getRedisTemplate(actualInstance);
    }

    /**
     * 根据数据实体获取对应的redis分片实例
     *
     * @param entity 数据实体
     * @return redis实例
     */
    public StringRedisTemplate getRedisByEntity(Object entity) {

        String logicTable = getEntityLogicTable(entity.getClass());

        // 未开启分片使用默认redis实例
        if (!enableSharding || !isSharding(logicTable)) {
            return getRedisByInstanceName(defaultRedisInstance);
        }
        Object shardingValue = idInjectorMap.get(logicTable).getShardingData(entity);
        return getRedisByShardingColumn(shardingValue, logicTable);
    }

    /**
     * 通过redis实例名获取对应的RedisTemplate用于操作对应的Redis
     *
     * @param redisInstanceName 分片结果
     * @return 分片结果映射的redis连接
     */
    public StringRedisTemplate getRedisByInstanceName(String redisInstanceName) {
        return redisShardingClient.getRedisTemplate(redisInstanceName);
    }

    /**
     * 获取缓存redis
     *
     * @return redis连接
     */
    public StringRedisTemplate getCacheRedis() {
        return redisShardingClient.getRedisTemplate(cacheRedisInstance);
    }

    /**
     * 获取实体的逻辑表名称
     *
     * @param entityCls 实体类
     * @return 逻辑表名称
     */
    public static String getEntityLogicTable(Class<?> entityCls) {

        // 如果不存在则put到集合，ConcurrentHashMap.computeIfAbsent线程安全
        return LOGIC_TABLES.computeIfAbsent(entityCls, cls -> {
            Table table = entityCls.getAnnotation(Table.class);
            if (null == table) {
                throw new RuntimeException(entityCls.getName() + "is not a entity class");
            }
            return table.name();
        });
    }

    /**
     * 表是否分片
     *
     * @param logicTable 逻辑表
     * @return true 表示分片 false不分片
     */
    public boolean isSharding(String logicTable) {

        return tableRules.containsKey(logicTable);
    }
}
